"""
Utility functions for multi-Gaussian fitting of lateral dose profiles.
横向 Dose 分布拟合, 使用 多高斯分布

The module includes:
- 多高斯模型 Multi-Gaussian model
- 优化的目标函数 Objective functions
- 多高斯模型参数估计 和 约束函数 Parameters and constraints
- Bridge class, for 模型参数Tuple 与 结构体命名字段 互相转换
- 工具函数: 归一化, 统计量等
"""

from typing import Any, Dict, Optional, Tuple

import numpy as np
from icecream import ic
from scipy.optimize import OptimizeResult, minimize

from .fitmodel import (
    BoundsIn,
    CnstrntsIn,
    FitConfig,
    GaussParams,
    MGaussIn,
    MidTailDtctCfg,
    MidTailDtctData,
    ObjctvArgs2,
    RunMGFitIn,
    RunMGFitRet,
    WtdStdDevIn,
    cvt_params_to_stctd,
)
from .fitmodel import myDftVals as dv
from .fitutils import mid_tail_detect_v2, mid_tail_weights, weighted_stddev


##00000000000000000000000000000000000000000
def multi_gauss_model(
    mIn: MGaussIn,
) -> np.ndarray:
    """
    N-Gaussian model function.

    Parameters:
        crds: Input coordinates
        params: Parameters array [normFactor, w2..wn, s1..sn, (c)]
            - normFactor: Overall normalization factor
            - w2..wn: Weights for 2nd to nth Gaussians (first weight is 1 - sum of others)
            - s1..sn: Sigmas for each Gaussian
            - c: Optional background constant
        n_gauss: Number of Gaussian components
        with_background: Whether to include background term

    Returns:
        Predicted y values according to the multi-Gaussian model
    """
    ## 使用 GaussParams 结构化对象来访问参数
    ## gauss_params = cvt_params_to_stctd(gIn.params, gIn.n_gauss, gIn.with_background)
    pars = mIn.gpars

    ## 归一化因子
    norm_factor = pars.norm_factor
    ## Create arrays for all weights and sigmas at once for vectorization
    all_weights = dv.fctr_symmetry * pars.all_weights
    all_sigmas = pars.sigmas

    ## Expand x to (n_gauss, len(x)) and compute all Gaussians at once
    x_expanded = mIn.crds[np.newaxis, :]  # Shape: (1, len(x))

    ## Compute all Gaussian components at once
    ## Shape: (n_gauss, len(x))
    gaussian_components = all_weights[:, np.newaxis] * np.exp(-0.5 * (x_expanded / all_sigmas[:, np.newaxis]) ** 2) / (np.sqrt(2 * np.pi) * all_sigmas[:, np.newaxis])

    ## Sum all components along the first axis (n_gauss dimension)
    y_pred = np.sum(gaussian_components, axis=0)  # Shape: (len(x),)

    ## Apply normalization factor
    y_pred *= norm_factor

    ## Add background if requested
    if pars.with_background and pars.background is not None:
        y_pred += pars.background

    return y_pred


def sigma_smooth_penalty(
    sigmas: np.ndarray,
    s_prev: Optional[np.ndarray],
    lam_s: float = dv.penalty_sigma_smooth,
) -> float:
    """
    Penalty term to ensure smooth transition of sigma values between adjacent layers.

    Parameters:
        s: Current sigma values
        s_prev: Previous sigma values (from adjacent layer)
        lam_s: Penalty strength

    Returns:
        Penalty value
    """
    if s_prev is None or len(sigmas) == 0 or len(s_prev) != len(sigmas):
        return 0.0

    # Compute relative difference and apply penalty using vectorized operations
    rel_diff = (sigmas - s_prev) / np.maximum(s_prev, dv.epsilon)
    return lam_s * np.sum(rel_diff**2)


def objective_log(
    gpars_ary: np.ndarray,
    oIn: ObjctvArgs2,
) -> float:
    """
    First stage objective function: log residual sum of squares with soft constraints.

    Parameters:
        params: Fitting parameters [normFactor, w2..wn, s1..sn, (c)]
        fitting_data: FittingData object with coordinates and values
        config: FittingConfig object with fitting parameters
        s_prev: Previous sigma values (for smoothness penalty)
        mid_tail_weight_para: Parameters for mid-tail weighting

    Returns:
        Objective function value
    """
    vals_tgt = oIn.vals_tgt  # 期望拟合到的目标
    fitCfg = oIn.fitcfg
    ngauss = fitCfg.ngauss
    lam_n = fitCfg.lam_n
    gpars = cvt_params_to_stctd(gpars_ary, ngauss)

    # All parameters must be non-negative and sigmas positive
    if np.any(gpars_ary < 0) or np.any(gpars.sigmas <= 0):
        return 1e20

    mgauss_in = MGaussIn(crds=oIn.crds, gpars=gpars)
    vals_p = multi_gauss_model(mgauss_in)  ## MGauss 预测值
    eps = dv.epsilon

    # Apply mid-tail weights if provided
    if oIn.mt_wts_pars is None:
        weights = 1.0
    else:
        # ic(oIn.mt_wts_pars)
        weights = mid_tail_weights(oIn.mt_wts_pars)

    ## Calculate log residual sum of squares (vectorized)
    vals_prdct = np.log(vals_p + eps)  # predict
    vals_tgt = np.log(vals_tgt + eps)
    residual = vals_prdct - vals_tgt
    loss = np.sum(weights * residual**2.0)

    # Soft constraint: normFactor should be close to 1
    loss += lam_n * np.pow(gpars.norm_factor - 1.0, 2.0)

    # Add smoothness penalty if previous sigmas are provided
    if oIn.sigma_prev is not None:
        loss += sigma_smooth_penalty(gpars.sigmas, oIn.sigma_prev)

    return loss


def objective_linear(
    gpars_ary: np.ndarray,
    oIn: ObjctvArgs2,
) -> float:
    """
    Second stage objective function: linear residual sum of squares with soft constraints.

    Parameters:
        params: Fitting parameters [normFactor, w2..wn, s1..sn, (c)]
        fitting_data: FittingData object with coordinates and values
        config: FittingConfig object with fitting parameters
        s_prev: Previous sigma values (for smoothness penalty)
        mid_tail_weight_para: Parameters for mid-tail weighting

    Returns:
        Objective function value
    """
    ## 必须输入 期望拟合到的目标 values
    vals_tgt = oIn.vals_tgt  # 期望拟合到的目标
    fitcfg = oIn.fitcfg
    ngauss = fitcfg.ngauss
    gpars = cvt_params_to_stctd(gpars_ary, ngauss)

    # All parameters must be non-negative and sigmas positive
    if np.any(gpars_ary < 0) or np.any(gpars.sigmas <= 0):
        return 1e20

    mgauss_in = MGaussIn(crds=oIn.crds, gpars=gpars)
    vals_p = multi_gauss_model(mgauss_in)  ## MGauss 预测值

    ## Apply mid-tail weights if provided
    if oIn.mt_wts_pars is None:
        weights = 1.0
    else:
        weights = mid_tail_weights(oIn.mt_wts_pars)

    ## Calculate linear residual sum of squares (vectorized)
    linear_residual = vals_p - vals_tgt
    loss = np.sum(weights * linear_residual**2.0)

    ## Soft constraint: normFactor should be close to 1
    loss += fitcfg.lam_n * np.pow(gpars.norm_factor - 1.0, 2.0)

    ## Add smoothness penalty if previous sigmas are provided
    if oIn.sigma_prev is not None:
        loss += sigma_smooth_penalty(gpars.sigmas, oIn.sigma_prev)

    return loss


def build_constraints(cIn: CnstrntsIn) -> list:
    """
    Build constraints for the optimization problem.

    Parameters:
        n_gauss: Number of Gaussian components
        is_post_bragg: Whether this is for post-Bragg peak region

    Returns:
        List of constraint dictionaries for scipy.optimize.minimize
    """
    constraints = []

    ## Weight sum constraint: sum of weights should be <= 1
    ## This means 1 - sum(w2..wn) >= w1 >= 0, so sum(w2..wn) <= 1
    constraints.append(
        {
            "type": "ineq",
            "fun": lambda gpars_ary: 1 - np.sum(gpars_ary[1 : cIn.ngauss]),
        }
    )

    ## Sigma ordering constraint: sigma values should be non-decreasing
    stt = cIn.ngauss  # 第 0 个 sigma
    end = cIn.ngauss + cIn.ngauss - 1  # 第 n-1 个 sigma
    for ix in range(stt, end):
        ## Skip first sigma constraint for post-Bragg if specified
        if cIn.is_post_bragg and ix == 0:
            continue
        constraints.append(
            {
                "type": "ineq",
                "fun": lambda gpars_ary: gpars_ary[ix + 1] - gpars_ary[ix],
            }
        )

    return constraints


def build_bounds(
    # crds: np.ndarray,
    # ngauss: int,
    bIn: BoundsIn,
) -> np.ndarray:
    """
    Build bounds for the optimization parameters.

    Parameters:
        crds: Input coordinates (used to determine max range for sigmas)
        n_gauss: Number of Gaussian components
        bound_pars:
        is_post_bragg: Whether this is for post-Bragg peak region

    Returns:
        List of (min, max) bounds for each parameter
    """
    # Use default optimization parameters if not provided
    ngauss = bIn.ngauss
    sig_prev = bIn.sig_prev

    ## Initial bounds [normFactor, w2..wn, s1..sn, (c)]

    sig_low = np.full(ngauss, dv.sigma_min)
    sig_upp = np.full(ngauss, bIn.crds.max() - bIn.crds.min())
    # ic(sig_low)
    # ic(sig_upp)

    if sig_prev is not None:
        sig_range: float = bIn.step_range_relax if bIn.is_post_bragg else bIn.step_range
        sig_low2 = (1.0 - sig_range) * sig_prev
        sig_low = np.fmax(sig_low, sig_low2)
        ## update sigmas upper bound
        sig_upp2 = (1.0 + sig_range) * sig_prev
        sig_upp = np.fmax(sig_low * 1.05, sig_upp2)

    lower = GaussParams(
        norm_factor=dv.norm_fctr_min,
        weights=np.full(ngauss - 1, dv.weight_min),
        sigmas=sig_low,
    ).to_array()

    upper = GaussParams(
        norm_factor=dv.norm_fctr_max,
        weights=np.full(ngauss - 1, 1.0),
        sigmas=sig_upp,
    ).to_array()

    bounds = np.column_stack((lower, upper))
    # ic(bounds)  ## Debug info
    return bounds


def fit_multi_gauss(
    # fitting_data: FitData,
    # config: Optional[FitConfig] = None,
    # init_params: Optional[np.ndarray] = None,
    # opt_params: Optional[OptimizeParams] = None,
    rIn: RunMGFitIn,
) -> RunMGFitRet:
    """
    Fit multi-Gaussian model to the data.

    Parameters:
        fitting_data: FittingData object with coordinates and values
        config: FittingConfig object with fitting parameters
        init_params: Initial parameter values
        opt_params: OptimizationParams object with optimization parameters

    Returns:
        Tuple of (fitted parameters, result object)
    """
    # Use default configurations if not provided
    fitCfg = rIn.fitcfg
    if fitCfg is None:
        fitCfg = FitConfig()

    # crds = fitting_data.coords
    # vals = fitting_data.values
    # n_gauss = config.ngauss
    # with_background = config.with_background
    # is_post_bragg = config.is_post_bragg
    # lam_n = config.lam_n
    # info_verbose = config.info_verbose
    # info_tag = config.info_tag
    # s_prev = opt_params.s_prev
    # crds = np.asarray(crds, dtype=float)
    # vals = np.asarray(vals, dtype=float)

    # Initialize parameters if not provided
    gpars_init: GaussParams
    if rIn.gpars_init is None:
        with_bkgrd = fitCfg.with_background
        bkgrd = rIn.vals_tgt.min() if with_bkgrd else None
        # Evenly distribute weights among all components except the first
        wts_init = np.ones(fitCfg.ngauss - 1) * 0.5 / (fitCfg.ngauss - 1)
        # Estimate initial sigma using weighted standard deviation
        sigma_guess = weighted_stddev(WtdStdDevIn(qty=rIn.crds, wts=rIn.vals_tgt))
        sigmas_init = np.array([sigma_guess * (0.8 + 2 * i) for i in range(fitCfg.ngauss)])
        # 使用 GaussParams 结构化对象来构建初始参数
        gpars_init = GaussParams(
            norm_factor=1.0,
            weights=wts_init,
            sigmas=sigmas_init,
            with_background=with_bkgrd,
            background=bkgrd,
        )
        if fitCfg.info_verbose:
            print("f[{info_tag}]Initial params: {gpars_init}")
    else:
        gpars_init = rIn.gpars_init

    ## Create data models for the optimization parameters
    # fitting_data = FitData(coords=crds, values=vals)
    # fitting_config = FitConfig(ngauss=n_gauss, with_background=with_background, lam_n=lam_n, is_post_bragg=is_post_bragg)
    ## Build bounds and constraints
    bounds = build_bounds(
        BoundsIn(
            ngauss=fitCfg.ngauss,
            crds=rIn.crds,
            is_post_bragg=fitCfg.is_post_bragg,
            sig_prev=rIn.sigma_prev,
        )
    )
    ##
    constraints = build_constraints(
        CnstrntsIn(
            ngauss=fitCfg.ngauss,
            is_post_bragg=fitCfg.is_post_bragg,
        )
    )

    ## First stage: preliminary fit to determine mid-tail weighting
    # ic(fitCfg)
    # print("type of fitCfg:", type(fitCfg))
    # print("type expected by ObjctvArgs2:", ObjctvArgs2.model_fields['fitcfg'].annotation)
    objArg2 = ObjctvArgs2(
        crds=rIn.crds,
        vals_tgt=rIn.vals_tgt,
        fitcfg=fitCfg,
        sigma_prev=rIn.sigma_prev,
        mt_wts_pars=None,
    )

    gpars_ary = gpars_init.to_array()
    ic(gpars_ary)
    ic(bounds)
    res0 = minimize(
        lambda gpars_ary: objective_log(gpars_ary, objArg2),
        gpars_ary,
        # args=None,
        method="SLSQP",
        bounds=bounds,
        constraints=constraints,
        options={"maxiter": 500, "ftol": 1e-10},
    )

    gpars_ret = cvt_params_to_stctd(
        res0.x,
        ngauss=fitCfg.ngauss,
        with_background=fitCfg.with_background,
    )
    vals_p0 = multi_gauss_model(MGaussIn(crds=rIn.crds, gpars=gpars_ret))  ## predict 0

    ## Detect mid-tail parameters for better fitting
    mt_wpars = mid_tail_detect_v2(
        MidTailDtctData(
            crds=rIn.crds,
            vals_tgt=rIn.vals_tgt,
            vals_p0=vals_p0,
        )
    )

    if fitCfg.info_verbose:
        prefix = f"[{fitCfg.info_tag}] " if fitCfg.info_tag is not None else ""
        print(f"{prefix}mid_tail_weight_para: center={mt_wpars.center:.3f}, width={{mt_wpars.width:.3f}}, lam_m={{mt_wpars.lam_m:.1f}}, n={{fitCfg.ngauss}}")

    ##========================================== Stage 1: Log residual fitting to emphasize tails
    objArg2 = ObjctvArgs2(
        crds=rIn.crds,
        vals_tgt=rIn.vals_tgt,
        fitcfg=fitCfg,
        sigma_prev=rIn.sigma_prev,
        mt_wts_pars=mt_wpars,
    )

    res1: OptimizeResult = minimize(
        lambda gpars_ary: objective_log(gpars_ary, objArg2),
        gpars_ary,  # res0.x,
        # args=None,
        method="SLSQP",
        bounds=bounds,
        constraints=constraints,
        options={"maxiter": 3000, "ftol": 1e-10},
    )

    ##========================================== Stage 2: Linear residual refinement
    res2 = minimize(
        lambda gpars_ary: objective_linear(gpars_ary, objArg2),
        res1.x,
        # args=None,
        method="SLSQP",
        bounds=bounds,
        constraints=constraints,
        options={"maxiter": 2000, "ftol": 1e-10},
    )

    if not res2.success:
        print("Optimization failed to converge: ", res2.message)

    gpars_out0 = cvt_params_to_stctd(res0.x, fitCfg.ngauss, with_background=fitCfg.with_background)
    gpars_out1 = cvt_params_to_stctd(res1.x, fitCfg.ngauss, with_background=fitCfg.with_background)
    gpars_out2 = cvt_params_to_stctd(res2.x, fitCfg.ngauss, with_background=fitCfg.with_background)

    return RunMGFitRet(
        gpars_out0,
        gpars_out1,
        gpars_out2,
        ret=res2,
    )
